from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma

from model.MyOllamaEmbeddings import MyOllamaEmbeddings

# 1. 初始化 Ollama 嵌入模型（确保 ollama serve 正在运行）
# embeddings = OllamaEmbeddings(
#     model="nomic-embed-text",  # 你部署的模型名
#     base_url="http://localhost:11434"  # Ollama 默认地址
# )

embeddings = MyOllamaEmbeddings(model="nomic-embed-text")

# 2. 准备一些示例文本
texts = [
    "我喜欢吃苹果",
    "机器学习很有趣",
    "Python是最好的编程语言",
    "今天天气很好",
    "人工智能正在改变世界"
]

# 3. 直接生成文本嵌入向量
print("生成文本嵌入...")
vectors = embeddings.embed_documents(texts)
print(f"生成的向量数量: {len(vectors)}") # 文本数量
print(f"每个向量的维度: {len(vectors[0])}") # 第一段文本的embeddings向量维度

# 4. 创建向量数据库
print("\n创建向量数据库...")
vectorstore = Chroma.from_texts(
    texts=texts,
    embedding=embeddings,
    persist_directory="./simple_db"  # 保存到本地
)

# 5. 进行相似性搜索

print("\n进行相似性搜索...")
query = "编程语言"
results = vectorstore.similarity_search(query, k=3)

# todo 这里输出结果不对，后期去查查
print(f"查询: '{query}'")
print("最相似的结果:")
for i, doc in enumerate(results):
    print(f"{i+1}. {doc.page_content}")
